import fitz
import os
from typing import Optional, List, Dict


class PDFLawReader:
    """
    基于PyMuPDF的PDF法律文档读取工具类
    支持读取文档信息、提取文本（全量/按章节）、保存文本等功能，适配法律文档结构化特点
    """

    def __init__(self, pdf_path: str):
        """
        初始化PDF读取器
        :param pdf_path: PDF文件路径（绝对路径或相对路径）
        :raises FileNotFoundError: 若指定路径的PDF文件不存在，抛出文件未找到异常
        :raises ValueError: 若指定路径文件不是有效的PDF，抛出值异常
        """
        # 验证文件是否存在
        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"PDF文件不存在：{pdf_path}")
        # 验证文件后缀是否为PDF
        if not pdf_path.lower().endswith(".pdf"):
            raise ValueError(f"文件不是有效的PDF格式：{pdf_path}")

        self.pdf_path = pdf_path
        self.doc: Optional[fitz.Document] = None  # PDF文档对象，初始化时为None
        self.total_pages: int = 0  # PDF总页数
        self.metadata: Dict[str, str] = {}  # PDF元数据（标题、作者、创建日期等）

    def open_document(self) -> None:
        """
        打开PDF文档，初始化文档对象、总页数和元数据
        :raises RuntimeError: 若文档打开失败，抛出运行时异常
        """
        try:
            self.doc = fitz.open(self.pdf_path)
            self.total_pages = len(self.doc)
            # 处理元数据：将PyMuPDF返回的元数据格式化为易读字典，空值替换为"未标注"
            raw_metadata = self.doc.metadata
            self.metadata = {
                "标题": raw_metadata.get("title", "未标注"),
                "作者": raw_metadata.get("author", "未标注"),
                "创建日期": raw_metadata.get("creationDate", "未标注"),
                "修改日期": raw_metadata.get("modDate", "未标注"),
                "页数": str(self.total_pages),
                "PDF版本": raw_metadata.get("format", "未标注").split("-")[-1] if raw_metadata.get("format") else "未标注"
            }
        except Exception as e:
            raise RuntimeError(f"PDF文档打开失败：{str(e)}") from e

    def close_document(self) -> None:
        """关闭PDF文档，释放资源（必须在操作完成后调用，避免内存泄漏）"""
        if self.doc is not None:
            self.doc.close()
            self.doc = None  # 重置文档对象为None，标记已关闭

    def get_document_info(self) -> Dict[str, str]:
        """
        获取PDF文档基本信息（需先调用open_document()打开文档）
        :return: 包含文档信息的字典，键包括：标题、作者、创建日期、修改日期、页数、PDF版本
        :raises RuntimeError: 若文档未打开，抛出运行时异常
        """
        if self.doc is None:
            raise RuntimeError("请先调用open_document()打开PDF文档")
        return self.metadata

    def extract_full_text(self) -> str:
        """
        提取PDF文档全部文本内容（需先调用open_document()打开文档）
        :return: 文档全文字符串（按页码顺序拼接，页间用空行分隔）
        :raises RuntimeError: 若文档未打开，抛出运行时异常
        """
        if self.doc is None:
            raise RuntimeError("请先调用open_document()打开PDF文档")

        full_text = []
        for page_num in range(self.total_pages):
            page = self.doc[page_num]
            # 提取当前页文本，PyMuPDF的get_text()会自动处理文本换行和排版
            page_text = page.get_text()
            # 为每一页文本添加页码标识（便于定位内容来源）
            full_text.append(f"=== 第{page_num + 1}页 ===")
            full_text.append(page_text.strip())  # 去除页首/页尾空行
        # 用两个空行分隔不同页面的文本，提升可读性
        return "\n\n".join(full_text)

    def extract_text_by_chapter(self, chapter_keywords: List[str]) -> Dict[str, str]:
        """
        按章节提取文本（适配法律文档的"第一章/第二章..."章节结构，需先打开文档）
        :param chapter_keywords: 章节关键词列表，如["第一章", "第二章", "第三章"]
        :return: 章节-文本字典，键为章节名（如"第一章 总则"），值为对应章节的完整文本
        :raises RuntimeError: 若文档未打开，抛出运行时异常
        :raises ValueError: 若章节关键词为空或格式无效，抛出值异常
        """
        if self.doc is None:
            raise RuntimeError("请先调用open_document()打开PDF文档")
        if not chapter_keywords or not all(isinstance(kw, str) for kw in chapter_keywords):
            raise ValueError("章节关键词列表不能为空，且需为字符串类型（如['第一章', '第二章']）")

        # 先获取全量文本，按段落分割（法律文档通常每段为独立语义单元）
        full_text = self.extract_full_text()
        paragraphs = [p.strip() for p in full_text.split("\n") if p.strip()]  # 过滤空段落

        chapter_text = {}
        current_chapter = None
        current_content = []

        for para in paragraphs:
            # 检查当前段落是否为章节标题（匹配"第一章 总则"这类格式）
            is_chapter_title = any(para.startswith(kw) for kw in chapter_keywords)
            if is_chapter_title:
                # 若已有正在收集的章节，先保存该章节内容
                if current_chapter is not None and current_content:
                    chapter_text[current_chapter] = "\n".join(current_content)
                # 重置当前章节和内容（更新为新章节）
                current_chapter = para
                current_content = [para]  # 章节标题本身也加入内容
            else:
                # 若已有当前章节，将段落加入该章节内容
                if current_chapter is not None:
                    current_content.append(para)

        # 保存最后一个章节的内容
        if current_chapter is not None and current_content:
            chapter_text[current_chapter] = "\n".join(current_content)

        return chapter_text

    def save_text_to_file(self, text: str, save_path: str, encoding: str = "utf-8") -> None:
        """
        将提取的文本保存到本地文件（支持txt格式）
        :param text: 待保存的文本内容
        :param save_path: 保存路径（需包含文件名，如"劳动合同法_全文.txt"）
        :param encoding: 文件编码，默认utf-8（适配中文）
        :raises IOError: 若保存路径无效或无写入权限，抛出IO异常
        """
        # 确保保存路径的目录存在，不存在则创建
        save_dir = os.path.dirname(save_path)
        if save_dir and not os.path.exists(save_dir):
            os.makedirs(save_dir, exist_ok=True)

        try:
            with open(save_path, "w", encoding=encoding) as f:
                f.write(text)
            print(f"文本已成功保存到：{os.path.abspath(save_path)}")
        except Exception as e:
            raise IOError(f"文本保存失败：{str(e)}") from e

    def __enter__(self):
        """上下文管理器入口：自动打开文档"""
        self.open_document()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口：自动关闭文档（无论是否发生异常）"""
        self.close_document()
        # 若发生异常，返回False让异常正常抛出；无异常返回True
        return exc_type is None
